import torch
from itertools import product as product
import numpy as np
from math import ceil

# <Class: PriorBox/>
class PriorBox(object):
    """
    Some information about PriorBox ~
    - min_sizes(list/tuple of list/tuple): 
    - steps(list): shows the feature maps size relative to image size below.
    - clip(boolean/int): If clip the prior boxes inside the image
    - image_size(list/tuple): Specified the input size of network
    """
    # <Method: __init__/>
    def __init__(
        self, 
        min_sizes=[[16, 32], [64, 128], [256, 512]], 
        steps=[8, 16, 32], 
        clip=False, 
        image_size=(512, 512)
        ):
        super(PriorBox, self).__init__()
        self._min_sizes = min_sizes
        self._steps = steps
        self._clip = clip
        self._image_size = image_size
        self._feature_maps = [[ceil(self._image_size[0]/step), ceil(self._image_size[1]/step)] for step in self._steps]
    # <Method: /__init__>
    
    # <Method: forward/>
    def forward(self):
        anchors = []
        for k, f in enumerate(self._feature_maps):
            min_sizes = self._min_sizes[k]
            for i, j in product(range(f[0]), range(f[1])):
                for min_size in min_sizes:
                    s_kx = min_size / self._image_size[1]
                    s_ky = min_size / self._image_size[0]
                    dense_cx = [x * self._steps[k] / self._image_size[1] for x in [j + 0.5]]
                    dense_cy = [y * self._steps[k] / self._image_size[0] for y in [i + 0.5]]
                    for cy, cx in product(dense_cy, dense_cx):
                        anchors += [cx, cy, s_kx, s_ky]
                    # end-for
                # end-for
            # end-for
        # end-for
        # back to torch land
        output = torch.Tensor(anchors).view(-1, 4)
        if self._clip:
            output.clamp_(max=1, min=0)
        # end-if
        return output
    # <Method: /forward>
# <Class: /PriorBox>